#include <cuda_runtime.h>
#include <cuda_fp16.h> 
#include "conv2d.h"

extern "C" __global__ void implicit_gemm_v0(mykernelParamType param)
{ 
    // 计算输出位置线性索引（NHW维度合并）
    int OhOw = blockIdx.x * blockDim.x + threadIdx.x;
    // 输出通道维度（K维度）
    int k = blockIdx.y * blockDim.y + threadIdx.y;
    // Batch维度
    int n = blockIdx.z;
    // 边界检查
    if(OhOw >= param.Oh*param.Ow || k >= param.k || n >= param.n)
        return;
    // 分解输出坐标
    int oh = OhOw / param.Ow;
    int ow = OhOw % param.Ow;
    int input_addr, weight_addr, output_addr;
    float sum = 0.0;
    output_addr = n * param.k * param.Oh * param.Ow + k * param.Oh * param.Ow + oh * param.Ow + ow;
    // 卷积核遍历
    for(int r = 0; r < param.r; r++) 
    {
        for(int s = 0; s < param.s; s++) 
        {
            int ih = oh * param.u - param.p + r;
            int iw = ow * param.v - param.q + s;
            if (ih >= 0 && ih < param.h && iw >= 0 && iw < param.w)
            {
                // 通道维度展开（隐式GEMM特性）
                for(int c = 0; c < param.c; c++)
                {
                    // 输入内存地址（NCHW布局）
                    input_addr = n * param.c * param.h * param.w + c * param.h * param.w + ih * param.w + iw;
                    // 权重内存地址（KCRS布局）
                    weight_addr = k * param.c * param.r * param.s + c * param.r * param.s + r * param.s + s;
                    sum += param.pin[input_addr] * param.pweight[weight_addr];
                }
            }
        }
    }
    param.pout[output_addr] = sum;
}
// 核函数启动配置
void launch_implicit_gemm_v0(unsigned int outh, unsigned int outw, unsigned int k, unsigned int n, mykernelParamType* param) 
{
    // 计算网格维度（向上取整）
    int blockx = (outh * outw + 15) / 16;  // 每个block处理16个输出位置
    int blocky = (k + 15) / 16;            // 每个block处理16个输出通道
    int blockz = n;                         // Batch维度直接映射
    dim3 block(16, 16, 1);  // 固定线程块尺寸
    dim3 grid(blockx, blocky, blockz);
    implicit_gemm_v0<<<grid, block>>>(*param);
}
